/* File: BaseFunctions.c */ 
/*
####????

*/

#include <R.h>
#include <Rinternals.h>
#include <Rmath.h>

//define the number of loci to use
enum{ n_disp_loci = 2 }; 
enum{ n_neutral_loci = 2 };

/*--------------------------structures----------------------------------------*/
/*define the structure of individual
	x & y are location & sex -- 0 is male & 1 is female
	Da & Db are alleles for dispersal linked loci
	Na & Nb are alleles for neutral loci */
typedef struct {
	int x, y;
	short sex;
	short Da[n_disp_loci], Db[n_disp_loci], Na[n_neutral_loci], Nb[n_neutral_loci];
} ind;






/*--------------------------global variables----------------------------------*/
ind * cur_ind = NULL; //current set of individuals
ind * new_ind = NULL; //next generation of individuals
int popsize; //global variable for population size
int spX, spY; //global for space size
int ** density = NULL; //the grid for taking density
int ** nmales = NULL; //the grid for taking nmales
int *** mInds = NULL; //the grid for taking indexes for males
int lambda, K; //globals for growth rate parameters


/*--------------------------function prototypes-------------------------------*/
void init_inds (int n);
int * disp_loc(int X, int Y);
double hass_comm(double lambda, int K, int N);
void disperse(float d_cost);
void grid_attributes(void);
void free_grid_memory(void);
int first_empty(int xx, int yy);

/*--------------------------function definition-------------------------------*/
/*Initialises a matrix with n individuals
	individuals have a location (spX, spY) and are 
	initilised with gene for p.disp
	based on Stable range edge functions */
void init_inds (int n) {
	cur_ind = Calloc(n, ind);//allocate memory for this set of individuals
	int ii, jj; //some iteration variables
	for (ii = 0; ii < n; ii++) {
		cur_ind[ii].x = round( unif_rand() * spX );	cur_ind[ii].y = round( unif_rand() * spY ); // get random x & x locations
		cur_ind[ii].sex = Rf_rbinom(1,0.5); //define the sex
		//printf("individual %i -- X = %i -- Y = %i -- sex = %i\nlinked loci = ",ii,cur_ind[ii].x,cur_ind[ii].y,cur_ind[ii].sex);
		for (jj = 0; jj < n_disp_loci; jj++) { //assign the initial allelles
			cur_ind[ii].Da[jj] = Rf_rbinom(1,0.5); cur_ind[ii].Db[jj] = Rf_rbinom(1,0.5); 
			//printf("%i%i",cur_ind[ii].Da[jj],cur_ind[ii].Db[jj]);
		} 
		//printf("\nneutral loci = ");
		for (jj = 0; jj < n_neutral_loci; jj++) { //assign the initial allelles
			cur_ind[ii].Na[jj] = Rf_rbinom(1,0.5); cur_ind[ii].Nb[jj] = Rf_rbinom(1,0.5);
			//printf("%i%i",cur_ind[ii].Na[jj],cur_ind[ii].Nb[jj]);			
		}
		//printf("\n\n");		
	}
	popsize = n; //define the current population size
}

/*calculates neighbours for an individual excluding nonsense neighbours and returns an array of new grid refs
	doesn't include source cell
	space is cylindrical */
int * disp_loc(int X, int Y){
	int ii; //some iteration variables
	int nX[8]={-1, 0, 1, -1, 1, -1, 0, 1}; //potential shifts in X and Y
	int nY[8]={1, 1, 1, 0, 0, -1, -1, -1};
	int XY[2]={-99, -99}; //array to take new location
	while (XY[1]<0 || XY[1]>spX){ //pull a new location within 0:spX
		ii=round(unif_rand() * 8);
		XY[0]=X+nX[ii]; XY[1] = Y+nY[ii];
	}
	if (XY[1]>spY) XY[1]=XY[1] % spY; //wrap Y axis
	if (XY[1]<0) XY[1]=XY[1] + spY;
	return XY;
}

/* hassel commins population growth where b=1 (contest competition)
	generates expected fecundity*/
double hass_comm(double lambda, int K, int N){
	return lambda/(1+(lambda-1)*N/K);
}

/* disperses all individuals with a survival cost*/
void disperse(float d_cost){
	int dloc[2]={-99, -99};
	for (int ii=0; ii<popsize; ii++){ //loop through all individuals
		double dprob=0, dcost=0;
		for (int jj=0; jj<n_disp_loci; jj++){
			dprob+=cur_ind[ii].Da[jj];
			dprob+=cur_ind[ii].Db[jj];
		}
		dprob=dprob/(2*n_disp_loci);
		dprob=Rf_rbinom(1, dprob);
		dcost=Rf_rbinom(1, d_cost);
		if (dprob==1 && dcost==0){
			int *tt = disp_loc(cur_ind[ii].x, cur_ind[ii].y);
			dloc[0] = tt[0]; dloc[1] = tt[1];
			cur_ind[ii].x=dloc[0];
			cur_ind[ii].y=dloc[1];
		}
		if (dprob==1 && dcost==1){ //marks dead individuals for removal later
			cur_ind[ii].sex=-1;
		}
	}
}

/* function to return the first empty postion in mInds vector */
int first_empty(int xx, int yy){ 
	for (int ii=0; ii<nmales[xx][yy]; ii++) {
		if (mInds[xx][yy][ii] == -99) { return(ii); }
	}
	return(-99);
}

/* populate grid information given a object of individuals */
void grid_attributes(void) {
	free_grid_memory();
	density = Calloc(spX,int *); nmales = Calloc(spX,int *); mInds = Calloc(spX,int **); //allocate memory for first dimension X
	for (int xx=0; xx<spX; xx++) { //initialise density and nmales matrices
		density[xx] = Calloc(spY,int); nmales[xx] = Calloc(spY,int); mInds[xx] = Calloc(spY,int*); //allocate memory for second dimension Y 
		for (int yy=0; yy<spY; yy++) { density[xx][yy]=0; nmales[xx][yy]=0; } //set everything to 0
	}
	for (int ii=0; ii<popsize; ii++) { //loop through all individuals & extract density info
		density[cur_ind[ii].x][cur_ind[ii].y]++; //add one to density
		if (cur_ind[ii].sex==0) { nmales[cur_ind[ii].x][cur_ind[ii].y]++; } //add one to males if sex == 0
	}	
	for (int xx=0; xx<spX; xx++) { //initialise density and nmales matrices
		for (int yy=0; yy<spY; yy++) {
			mInds[xx][yy] = Calloc(nmales[xx][yy],int); //allocate 3rd dim of this matrix
			for (int zz=0; zz<nmales[xx][yy]; zz++) { mInds[xx][yy][zz] = -99; }  
		}
	}
	for (int ii=0; ii<popsize; ii++) { //loop through all individuals & get the male index
		if (cur_ind[ii].sex==0) { 
			int xx = cur_ind[ii].x; int yy = cur_ind[ii].y; //get the position
			mInds[xx][yy][first_empty(xx,yy)] = ii;  //add one to males if sex == 0
		}
	}	
}

/* function to free up grid info from memory */
void free_grid_memory(void){
	if (density != NULL) { 	
		for (int xx=0; xx<spX; xx++) { //initialise density and nmales matrices
			for (int yy=0; yy<spY; yy++) { Free(mInds[xx][yy]); }
			Free(density[xx]); Free(nmales[xx]); Free(mInds[xx]);
		}
		Free(density); Free(nmales); Free(mInds);
	}
}

/* reproduces individuals based on lambda and grid cell K */

void reproduce(void){
	int nn=0; //counter for new_ind
	for (int ii=0; ii<popsize; ii++){
		if (cur_ind[ii].sex==1){ //females only
			double E = hass_comm(lambda, K, density[cur_ind[ii].x][cur_ind[ii].y]); //expected offspring
			int n_offsp = Rf_rpois(E); //actual offspring
			new_ind = Realloc(nn+n_offsp, ind);//allocate memory for this set of individuals CHECK THIS!!
			int male = round( unif_rand() * nmales[cur_ind[ii].x][cur_ind[ii].y] ); //choose male
			male = mInds[cur_ind[ii].x][cur_ind[ii].y][male];
			//Could probably do the following loop more efficiently with: ind temp = cur_ind[ii], but I got scared 
			for (int jj=0; jj<n_offsp; jj++;){ //step through offspring and allocate to new_ind
				new_ind[nn].x=cur_ind[ii].x;
				new_ind[nn].y=cur_ind[ii].y;
				new_ind[nn].sex=Rf_rbinom(1, 0.5);
				for (int ll=0; ll<n_disp_loci; ll++){ //assign dispersal alleles from both parents
					if (Rf_rbinom(1, 0.5)==1) {
						new_ind[nn].Da[ll]=cur_ind[male].Da[ll];
						new_ind[nn].Db[ll]=cur_ind[ii].Db[ll];
					}
					else {
						new_ind[nn].Da[ll]=cur_ind[ii].Da[ll];
						new_ind[nn].Db[ll]=cur_ind[male].Db[ll];
					}
				}
				for (int ll=0; ll<n_neutral_loci; ll++){ //assign neutral alleles from both parents
					if (Rf_rbinom(1, 0.5)==1) {
						new_ind[nn].Na[ll]=cur_ind[male].Na[ll];
						new_ind[nn].Nb[ll]=cur_ind[ii].Nb[ll];
					}
					else {
						new_ind[nn].Na[ll]=cur_ind[ii].Na[ll];
						new_ind[nn].Nb[ll]=cur_ind[male].Nb[ll];
					}
				}
				nn++; //increment new_ind counter
			}
		}
	}
	popsize=nn; // adjust popsize
	&cur_ind=&new_ind; //pass pointer JEREMY CHECK THIS!!
	new_ind = Realloc(0, ind);//allocate zero memory for this set of individuals
}




/* define the mother function
	R_n is the starting number of individuals
	R_spX & R_spY are the dimensions of the landscape
	R_d_cost is the dispersal cost */
SEXP mother(SEXP R_n, SEXP R_spX, SEXP R_spY, SEXP R_d_cost, SEXP R_lambda, SEXP R_K) {
	GetRNGstate();
	spX = INTEGER(coerceVector(R_spX,INTSXP))[0]; //size of the habitat matrix - X
	spY = INTEGER(coerceVector(R_spY,INTSXP))[0]; //size of the habitat matrix - Y
	
	float d_cost = REAL(coerceVector(R_d_cost,REALSXP))[0]; //dispersal cost
	lambda = INTEGER(coerceVector(R_lambda,INTSXP))[0]; //lambda
	K = INTEGER(coerceVector(R_K,INTSXP))[0]; //K
	
	int ii, jj; //some iteration variables	
	init_inds(INTEGER(coerceVector(R_n,INTSXP))[0]); //populate the starting set of individuals
	grid_attributes();//get the grid attributes
	
	//return some useless info
	SEXP res; PROTECT(res = allocVector(INTSXP,1));
	INTEGER(res)[0] = -9999;

	PutRNGstate();
	//release and free up memory
	Free(cur_ind);
	free_grid_memory();
	UNPROTECT(1);
    return res;
}

